

SPECIAL_TOKENS_MAP = {
    "qwen": {
        "global_start": "",
        "role_start": "<|im_start|>",
        "role_end": "",
        "between_role_and_content": "\n",
        "eot": "<|im_end|>",
        "between_eot_and_role_start": "\n", 
    },
    "llama3": {
        "global_start": "<|begin_of_text|>",
        "role_start": "<|start_header_id|>",
        "role_end": "<|end_header_id|>",
        "between_role_and_content": "\n\n",
        "eot": "<|eot_id|>",
        "between_eot_and_role_start": "", 
    },
    "phi3": {
        "global_start": "",
        "role_start": "",
        "role_end": "",
        "between_role_and_content": "",
        "eot": "<|end|>",
        "between_eot_and_role_start": "", 
    },
    # "megrez": {
    #     "global_start": "",
    #     "role_start": "<|role_start|>",
    #     "role_end": "<|role_end|>",
    #     "between_role_and_content": "",
    #     "eot": "<|eos|>",
    #     "between_eot_and_role_start": "", 
    # },
    "megrez": {
        "global_start": "",
        "role_start": "<|role_start|>",
        "role_end": "<|role_end|>",
        "between_role_and_content": "",
        "eot": "<|turn_end|>",
        "between_eot_and_role_start": "", 
    },
}


ROLES_MAP = {
    "qwen": {
        "SYSTEM": "system",
        "USER": "user",
        "ASSISTANT": "assistant",
        "FUNCTION RESPONSE": "observation"
    },
    "llama3": {
        "SYSTEM": "system",
        "USER": "user",
        "ASSISTANT": "assistant",
        "FUNCTION RESPONSE": "observation"
    },
    "phi3": {
        "SYSTEM": "<|system|>",
        "USER": "<|user|>",
        "ASSISTANT": "<|assistant|>",
        "FUNCTION RESPONSE": "<|observation|>"
    },
    "megrez": {
        "SYSTEM": "system",
        "USER": "user",
        "ASSISTANT": "assistant",
        "FUNCTION RESPONSE": "observation"
    }
}